# -*- coding: utf-8 -*-
import json
import os
import datetime
import time
from mongoengine import *
import copy
import slugify
#这里的RetryDocument覆盖mongoengine默认的Document,支持mongodb连接失败无限重试
#如果不想重试可以注释掉此行
from apptrack.backend.mongo_retry import RetryDocument as Document
from retry import retry
from pymongo.errors import AutoReconnect, OperationFailure
from ..utils import get_boolean

def build_expired_index(document_class,key,timestamp):
    '''
        设置过期索引
    '''
    index = {
        'fields': [key],
        'expireAfterSeconds':timestamp
    }
    document_class._meta['indexes'].append(index)
    index_specs = document_class._build_index_specs([index])
    document_class._meta['index_specs'].extend(index_specs) 
        
def init(**kwargs):
    expired_time = kwargs.pop('expired',None)
    #如果设置过期时间,设置过期索引,否则不设置
    if expired_time:
        build_expired_index(Span,'created_at',expired_time)
        build_expired_index(Context,'created_at',expired_time)
        build_expired_index(Tag,'created_at',expired_time)
        build_expired_index(Ring,'created_at',expired_time)
        build_expired_index(Log,'timestamp',expired_time)
        
    for k, v in kwargs.items():
        conn_params = copy.deepcopy(v)
        name = conn_params.pop('name')
        register_connection(k, name, **conn_params)
        
class Trace(Document):
    name = StringField(max_length=200)
    trace_id = StringField()
    created_at = DateTimeField(required=True,default=datetime.datetime.utcnow)
    #追踪的结束时间,如果是全局唯一的追踪实例,则没有结束时间,否则应该有
    end_at = DateTimeField()
    
    meta = {
        'db_alias':'default'
    }
    
    @staticmethod
    def get_trace(trace_id):
        trace = Trace.objects(trace_id=trace_id).first()
        return trace
        
    @classmethod
    def create_trace(cls,trace_id,service_name):
        if cls.get_trace(trace_id):
            return 
        Trace(name=service_name,trace_id=trace_id).save()
        
class Span(Document):
    start_time = DateTimeField()
    end_time = DateTimeField()
    name = StringField(max_length=200)
    parent_id = ObjectIdField()
    created_at = DateTimeField(required=True,default=datetime.datetime.utcnow)
    
    meta = {
        'db_alias':'default'
    }
    
    @staticmethod  
    def get_span(span_id):
        span = Span.objects(id=span_id).first()
        if span is None:
            raise RuntimeError("span id %s is not exist"%span_id)
        return span
        
    @staticmethod
    def make_mongo_dict(kwargs):
        d = {}
        for k,v in kwargs.items():
            d[slugify.slugify(k, separator='_')] = v
        return d
    

def create_new_id():
    '''
    '''
    span = Span()
    span.save()
    return str(span.id)
    
@retry(AutoReconnect, tries=4000, delay=10)
@retry(OperationFailure, tries=4000, delay=10)
def put_tag(span_id,key,value):
    '''
    '''
    v_type = type(value).__name__
    if v_type == "unicode":
        v_type = "str"
    if v_type == "bool":
        v = str(value).lower()
    else:
        v = str(value) 
    tag = Tag(span_id=span_id,name=key,value=v,v_type=v_type)
    tag.save()
    
def put_log(span_id,level,msg,rings = [],**kwargs):
    '''
    '''
    log = Log()
    Span.get_span(span_id)
    log.span_id = span_id
    log.event = level
    log.message = msg
    log.payload = Span.make_mongo_dict(kwargs)
    log.save()
    for ring in rings:
        Ring(log_id=log.id,name=ring).save()

def finish_span(span_id):
    '''
    '''
    span = Span.get_span(span_id)
    span.end_time = datetime.datetime.utcnow()
    span.save()
    
def start_span(span_id,parent_id=None,tags=[],context=None):
    '''
    '''
    span = Span.get_span(span_id)
    span.start_time = datetime.datetime.utcnow()
    if parent_id is not None:
        Span.get_span(parent_id)
        span.parent_id = parent_id
    span.save()
    for k,v in tags:
        put_tag(span_id,k,v)
        
    if context is not None:
        context['baggage'] = Span.make_mongo_dict(context['baggage'])
        Context(**context).save()
    

def set_span_name(span_id,name):
    '''
    '''
    span = Span.get_span(span_id)
    span.name = name
    span.save()

class Log(Document):
    span_id = ObjectIdField(required=True)
    event = StringField(max_length=20)
    timestamp = DateTimeField(default=datetime.datetime.utcnow)
    message = StringField()
    #这个字段要求必须定义'default'数据库别名
    payload = DictField()
    
    meta = {
        'db_alias':'default'
    }

class Tag(Document):
    span_id = ObjectIdField(required=True)
    name = StringField()
    value = StringField()
    v_type = StringField()
    created_at = DateTimeField(required=True,default=datetime.datetime.utcnow)
    
    meta = {
        'db_alias':'default'
    }

class Context(Document):
    span_id = ObjectIdField(required=True)
    parent_id = ObjectIdField()
    baggage = DictField()
    flags = IntField()
    trace_id = StringField(max_length=100)
    created_at = DateTimeField(required=True,default=datetime.datetime.utcnow)
    
    meta = {
        'db_alias':'default'
    }
    

class Ring(Document):
    log_id = ObjectIdField(required=True)
    name = StringField(max_length=20)
    created_at = DateTimeField(required=True,default=datetime.datetime.utcnow)
    meta = {
        'db_alias':'default'
    }
    
def obj_to_dict(obj):
    return dict(obj.to_mongo())
    
def to_list_dict(objs,**key_values):
    lst = []
    for obj in objs:
        data = obj_to_dict(obj)
        if key_values:
            data[key_values['key']]=key_values['value']
        lst.append(data)
    return lst

def full_span_id(span_id):
    parent_id_list = [str(span_id)]
    parent_id = span_id
    while parent_id:
        span = Span.objects(id=parent_id).first()
        parent_id = span.parent_id
        if parent_id:
            parent_id_list.append(str(parent_id))
    parent_id_list.reverse()
    return ".".join(parent_id_list)

def get_span_info(span_id,append_tag=True,recursive=False,to_dict=True):
    span = Span.get_span(span_id)
    if to_dict:
        span_data = obj_to_dict(span)
        span_data['id'] = full_span_id(span_id)
        if append_tag:
            tags = Tag.objects(span_id=span_id)
            tag_list = to_list_dict(tags)
            span_data['tags'] = tag_list
        span_data['childs'] = []
        if recursive:
            child_spans = Span.objects(parent_id=span_id)
            childs = []
            for child_span in child_spans:
                childs.append(get_span_info(child_span.id,append_tag,recursive,to_dict))
            span_data['childs'] = childs
            
        context = Context.objects(span_id=span_id).first()
        if not context:
            span_data['context'] = None
        else:
            trace_id = context.trace_id
            trace = Trace.get_trace(trace_id)
            span_data['context'] = obj_to_dict(context)
            span_data['context']['trace_info'] = obj_to_dict(trace)
        return span_data
    return span

def get_span_logs(span_id,timestamp=None,log_ring=[],recursive=False,log_list = []):
    q = {
        'span_id':span_id
    }
    if timestamp is not None:
        q["timestamp__gt"] = timestamp
    logs = Log.objects(**q)
    span = Span.get_span(span_id)
    span_name = span.name
    #过滤符合日志环的日志,其它的日志不需要
    if log_ring:
        ring_logs = []
        for log in logs:
            append_log=False
            for ring in log_ring:
                if Ring.objects(name=ring,log_id=log.id).first():
                    append_log = True
                    break
            if append_log:
                ring_logs.append(log)
        logs = ring_logs
    
    log_list.extend(to_list_dict(logs,**{'key':'span_name','value':span_name}))
    if recursive:
        child_span_ids = Span.objects(parent_id=span_id).only("id")
        for child_span_id in child_span_ids:
            get_span_logs(child_span_id.id,timestamp,log_ring,recursive,log_list)

def form_span_id(span_id):
    obj_span_id = hex(span_id).replace("0x","").replace("L","")
    return obj_span_id
    
def update_context(context):
    obj = Context.objects(span_id=context['span_id']).first()
    obj.baggage = Span.make_mongo_dict(context['baggage'])
    obj.flags = context['flags']
    obj.save()
    
def get_span_context(span_id):
    context = Context.objects(span_id=span_id).first()
    if not context:
        return None
    return obj_to_dict(context)
    
def get_tag(span_id,key,default):
    tag = Tag.objects(span_id=span_id,name=key).order_by('-created_at').first()
    if tag:
        if tag.v_type == "bool":
            return get_boolean(tag.value, default)
        elif tag.v_type == "int":
            return int(tag.value)
        elif tag.v_type in ['tuple','list','dict']:
            return eval(tag.value)
        else:
            return tag.value
    return default
    
def create_trace(trace_id,service_name):
    Trace.create_trace(trace_id,service_name)